import random

import os

import tensorflow
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.layers import Convolution2D, MaxPooling2D
from tensorflow.keras.models import load_model
from tensorflow.keras import backend as K

from load_dataset import load_dataset, resize_image


image_size = 64

class Dataset:
    def __init__(self, path_name):
        # 训练集
        self.train_images = None
        self.train_labels = None

        # 测试集
        self.test_images = None
        self.test_labels = None

        # 数据集加载路径
        self.path_name = path_name
        # 图像种类
        self.user_num = len(os.listdir(path_name))
        #当前库采用的维度顺序
        self.input_shape = None

        # 加载数据集并按照交叉验证的原则划分数据集并进行相关预处理工作
    def load(self, img_rows=image_size, img_cols=image_size,
            img_channels=3):
        # 数据种类
        nb_classes = self.user_num
        # 加载数据集到内存
        images, labels = load_dataset(self.path_name)

        train_images, test_images, train_labels, test_labels= train_test_split(images, labels, test_size=0.3,
                                                       random_state=random.randint(0,100))

        # 当前的维度顺序如果为'th'，则输入图片数据时的顺序为：channels,rows,cols，否则:rows,cols,channels
        # 这部分代码就是根据keras库要求的维度顺序重组训练数据集
        if K.image_data_format() == 'th':
            train_images = train_images.reshape(train_images.shape[0], img_channels, img_rows, img_cols)
            test_images = test_images.reshape(test_images.shape[0], img_channels, img_rows, img_cols)
            self.input_shape = (img_channels, img_rows, img_cols)
        else:
            train_images = train_images.reshape(train_images.shape[0], img_rows, img_cols, img_channels)
            test_images = test_images.reshape(test_images.shape[0], img_rows, img_cols, img_channels)
            self.input_shape = (img_rows, img_cols, img_channels)

            # 输出训练集、验证集、测试集的数量
            print(train_images.shape[0], 'train samples')
            print(test_images.shape[0], 'test samples')

            # 模型使用categorical_crossentropy作为损失函数，因此需要根据类别数量nb_classes将
            # 类别标签进行one-hot编码使其向量化，在这里我们的类别有六种，经过转化后标签数据变为六维
            train_labels = tensorflow.one_hot(train_labels, nb_classes)
            test_labels = tensorflow.one_hot(test_labels, nb_classes)

            # 像素数据浮点化以便归一化
            train_images = train_images.astype('float32')
            test_images = test_images.astype('float32')

            # 将其归一化,图像的各像素值归一化到0~1区间
            train_images /= 255
            test_images /= 255

            self.train_images = train_images
            self.test_images = test_images
            self.train_labels = train_labels
            self.test_labels = test_labels


# CNN网络模型类
class Model:
    def __init__(self):
        self.model = None

        # 建立模型
    def build_model(self, dataset,nb_classes=6):

        # 构建一个空的网络模型，它是一个线性堆叠模型，各神经网络层会被顺序添加，专业名称为序贯模型或线性堆叠模型
        self.model = Sequential()

        self.model.add(Convolution2D(filters=6,
                                      kernel_size=(5, 5),
                                      strides=(1, 1),
                                      activation='relu',
                                      input_shape=dataset.input_shape))
        self.model.add(MaxPooling2D())
        self.model.add(Convolution2D(filters=16,
                                      kernel_size=(5, 5),
                                      strides=(1, 1),
                                      activation='relu'))
        self.model.add(MaxPooling2D())
        self.model.add(Flatten())
        self.model.add(Dense(120, activation='relu'))
        self.model.add(Dense(84, activation='relu'))
        self.model.add(Dense(nb_classes, activation='softmax'))

        # 输出模型概况
        self.model.summary()

        # 训练模型
    def train(self, dataset, batch_size = 64, nb_epoch = 20):
        self.model.compile(loss='categorical_crossentropy',
                            optimizer='adam',
                            metrics='acc')  # 完成实际的模型配置工作

        self.model.fit(dataset.train_images,
                        dataset.train_labels,
                        batch_size = batch_size,
                        epochs = nb_epoch,
                        validation_split=0.2,
                       )

    MODEL_PATH = './model/face.model.h5'

    def save_model(self, file_path=MODEL_PATH):
        self.model.save(file_path)

    def load_model(self, file_path=MODEL_PATH):
        self.model = load_model(file_path)

    def evaluate(self, dataset):
        score = self.model.evaluate(dataset.test_images, dataset.test_labels, verbose=1)
        print("%s: %.2f%%" % (self.model.metrics_names[1], score[1] * 100))

    # 识别人脸
    def predict(self, image):
        # 依然是根据后端系统确定维度顺序
        if K.image_data_format() == 'th' and image.shape != (1, 3, image_size, image_size):
            image = resize_image(image)  # 尺寸必须与训练集一致都应该是IMAGE_SIZE x IMAGE_SIZE
            image = image.reshape((1, 3, image_size, image_size))  # 与模型训练不同，这次只是针对1张图片进行预测
        elif K.image_data_format() == 'tf' and image.shape != (1, image_size, image_size, 3):
            image = resize_image(image)
            image = image.reshape((1, image_size, image_size, 3))

        # 浮点并归一化
        image = image.astype('float32')
        image /= 255

        # 给出输入属于各个类别的概率
        result_probability = self.model.predict_proba(image)

        # 给出类别预测
        if max(result_probability[0]) >= 0.9:
            result = self.model.predict_classes(image)
            print('result:', result)
            # 返回类别预测结果
            return result[0]
        else:
            return -1


if __name__ == '__main__':
    user_num = len(os.listdir('./all_photo/'))

    dataset = Dataset('./all_photo/')
    dataset.load()

    model = Model()
    model.build_model(dataset, nb_classes=user_num)

    # 先前添加的测试build_model()函数的代码
    model.build_model(dataset, nb_classes=user_num)
    # 测试训练函数的代码
    model.train(dataset)

    model.save_model(file_path='./model/face.model.h5')
